import json
import random
from collections import OrderedDict

raw_data_path = 'F:\dataset\doctor\\train_data.json'

SYSTEM = """你是一名眼科专家，你需要解答患者的疑问，提供准确的回答，必要时，提醒患者及时挂号就医。
当患者对症状描述不清时，你需要循序渐进的引导患者，询问患者的症状，以便给出准确的诊断。
当患者上传眼底图时，你需调用Agent判断眼底图中是否存在青光眼和糖尿病视网膜病变，若存在，需要提醒患者及时就医。
"""


# 判断数据是不是与眼科相关的数据
def judge_eye_data(conversation):
    keywords = [
        '眼', '视力', '视网膜', '青光眼', '白内障', '糖尿病视网膜病变', '糖尿病', '视神经', '视野', 'OCT',
        'OCTA', 'fundus', 'eye', '网脱'
    ]
    for item in conversation:
        for keyword in keywords:
            if keyword in (item['input'] + item['output']):
                return True
    return False


# 处理原始数据,并保存
def process_data(data_path):
    # 读取原始数据
    with open(data_path, 'r', encoding='utf-8') as file:
        data = json.load(file)

    # 最终输出
    final_res = []

    for item in data:
        tmp = {'conversation': []}
        if len(item) % 2 == 0:
            for i in range(0,len(item),2):
                tmp['conversation'].append({
                    'input': item[i].lstrip('病人：'),
                    'output': item[i + 1].lstrip('医生：')
                })
            tmp['conversation'][0]['system'] = SYSTEM
            flag = judge_eye_data(tmp['conversation'])
            if flag:
                final_res.append(tmp)
    with open('data/process_data_eye.json','w',encoding='utf-8') as file:
        json.dump(final_res, file, ensure_ascii=False, indent=4)
    print(len(final_res))
    print(final_res[0])


# 拆分数据集
def split_data(process_data_path):
    with open(process_data_path, 'r', encoding='utf-8') as file:
        data = json.load(file)

    n = len(data) // 30
    sampled_data = random.sample(data, n)
    with open('data/split_eye_data.json', 'w', encoding='utf-8') as file:
        json.dump(sampled_data, file, indent=4,ensure_ascii=False)
    print(len(sampled_data))

split_data('data/process_data_eye.json')





